import secrets
import rsa
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
import backend.DID.Key_Management as Key_Management
import backend.DID.BlockChain as BlockChain
import backend.DID.DataBase as DataBase
import backend.DID.En_Decrypt_Management as En_Decrypt_Management 
import socket


socket_conn = None
shared_key = None
private_key = None  # 本地私钥
public_key = None   # 本地公钥
receiver_did = None
receiver_public_key = None
pending_conn = None  # 存储已完成密钥协商的连接
server_socket = None        # 保存监听 socket


def search_publickey(DID):
    id_ = {
        'id': f'{DID}'
    }
    table = 'blockchain'
    search_data = 'publickey'

    publickey_bytes = BlockChain.query_data(table, id_, search_data).encode('utf-8')
    publickey = rsa.PublicKey.load_pkcs1(publickey_bytes)
    return publickey


def search_privatekey(DID):
    return DataBase.private_key_changes(DID)



def recvall(sock, length):                          # 按指定字节数完整接收

    data = b''
    while len(data) < length:
        part = sock.recv(length - len(data))
        if not part:
            raise ConnectionError("socket closed")
        data += part
    return data

def send_with_len(sock, data):                      # 前置 4 字节长度后发送
    sock.sendall(len(data).to_bytes(4, 'big') + data)

def recv_with_len(sock):                            # 先读 4 字节长度，再读正文
    header = recvall(sock, 4)
    body_len = int.from_bytes(header, 'big')
    return recvall(sock, body_len)






def DID_listen(own_did, port, host='0.0.0.0'):
    global socket_conn, shared_key, receiver_did, private_key, public_key, receiver_public_key, pending_conn,server_socket

    private_key = search_privatekey(own_did)

    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)     # 建立TCP
    s.bind((host, port))
    s.listen(1)
    server_socket=s
    print(f"在端口{port}监听连接...")

    conn, addr = s.accept()
    print(f"收到来自 {addr} 的连接请求")

    try:
        # 互换 DID
        receiver_did = recv_with_len(conn).decode()
        print(f"接收方DID：{receiver_did}")
        send_with_len(conn, own_did.encode())

        # 接收密文（共享密钥 + random_a）
        receiver_public_key = search_publickey(receiver_did)
        encrypt_data = recv_with_len(conn)
        decrypt_data = En_Decrypt_Management.decrypt_by_privatekey(encrypt_data, private_key)
        if isinstance(decrypt_data, str):
            decrypt_data = decrypt_data.encode('utf-8')

        key_len = 16 if len(decrypt_data) == 32 else len(decrypt_data) - 16
        shared_key = decrypt_data[:key_len]
        random_a = decrypt_data[key_len:key_len + 16]

        # 生成 random_b 并回包
        random_b = secrets.token_bytes(16)
        data_to_send = random_a + random_b
        encrypt_response = En_Decrypt_Management.encrypt_by_sharekey(shared_key, data_to_send)
        send_with_len(conn, encrypt_response)

        # 验证 random_b
        encrypted_b = recv_with_len(conn)
        decrypted_b = En_Decrypt_Management.decrypt_by_sharekey(shared_key, encrypted_b)
        if isinstance(decrypted_b, str):
            decrypted_b = decrypted_b.encode('utf-8')

        if decrypted_b == random_b:
            print("密钥协商成功")
            print(f"共享密钥: {shared_key.hex()}")
            pending_conn = conn
            socket_conn = conn
            return True
        else:
            print("密钥协商失败，随机数不同")
            conn.close()
            return False

    except Exception as e:
        print(f"DID密钥协商失败：{e}")
        conn.close()
        return False


def DID_connect(receiver_did, port, sender_did, host):
    global socket_conn, shared_key, private_key, receiver_public_key

    private_key = search_privatekey(sender_did)

    socket_conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    socket_conn.connect((host, port))
    print(f"连接到 {host}:{port}，对方DID:{receiver_did} ")

    try:
        #  互换 DID
        send_with_len(socket_conn, sender_did.encode())
        received_did = recv_with_len(socket_conn).decode()
        print(f"收到对方DID: {received_did}")
        if receiver_did != received_did:
            print(f"DID不匹配: 期望{receiver_did}, 实际{received_did}")
            return False

        # 发送（共享密钥 + random_a）的密文
        receiver_public_key = search_publickey(receiver_did)
        shared_key = Key_Management.generate_aes_key()
        random_a = secrets.token_bytes(16)
        data_to_encrypt = shared_key + random_a
        encrypted_data = En_Decrypt_Management.encrypt_by_publickey(data_to_encrypt, receiver_public_key)
        send_with_len(socket_conn, encrypted_data)

        # 接收 random_a + random_b
        encrypted_response = recv_with_len(socket_conn)
        decrypted_response = En_Decrypt_Management.decrypt_by_sharekey(shared_key, encrypted_response)
        if isinstance(decrypted_response, str):
            decrypted_response = decrypted_response.encode('utf-8')

        key_len = len(shared_key)
        received_a = decrypted_response[:16]                  # 这里 random_a 固定 16
        random_b = decrypted_response[16:32]

        if received_a != random_a:
            print("密钥协商失败: 随机数A验证不通过")
            return False

        # 回传 random_b
        encrypted_b = En_Decrypt_Management.encrypt_by_sharekey(shared_key, random_b)
        send_with_len(socket_conn, encrypted_b)
        print("密钥协商成功")
        print(f"共享密钥: {shared_key.hex()}")
        return True
    except Exception as e:
        print(f"DID密钥协商失败：{e}")
        socket_conn.close()
        return False


def DID_send(data):
    global socket_conn, shared_key
    try:
        if not socket_conn or not shared_key:
            print("DID发送数据失败: 未建立有效连接或无共享密钥")
            return False
        socket_conn.getpeername()

        if isinstance(data, str):
            data = data.encode('utf-8')

        # print(f"datadata:{data}")
        encrypted_data = En_Decrypt_Management.encrypt_by_sharekey(shared_key, data)
        # print(f"send中：{encrypted_data}")
        send_with_len(socket_conn, encrypted_data)
        print("发送成功")
        return True
    except Exception as e:
        print(f"DID发送数据失败: {e}")
        return False


def DID_recv():
    global socket_conn, shared_key
    try:
        if not socket_conn or not shared_key:
            print("DID接收数据失败: 未建立有效连接或无共享密钥")
            return None
        socket_conn.getpeername()

        encrypted_data = recv_with_len(socket_conn)
        if not encrypted_data:
            print("未接受到数据")
            return None
        return En_Decrypt_Management.decrypt_by_sharekey(shared_key, encrypted_data)
    except Exception as e:
        print(f"DID接收数据失败: {e}")
        return None



def DID_close():
    global socket_conn, shared_key, private_key, receiver_did, receiver_public_key, server_socket, pending_conn

    try:
        if socket_conn:
            socket_conn.close()
            socket_conn = None

        if server_socket:  # ：关闭服务器套接字
            server_socket.close()
            server_socket = None

        if pending_conn and pending_conn != socket_conn:  # ：关闭待处理连接
            pending_conn.close()

        # 重置所有全局变量
        shared_key = None
        private_key = None
        receiver_did = None
        receiver_public_key = None
        pending_conn = None
        print("连接已关闭，密钥已销毁")

    except Exception as e:
        print(f"DID通信关闭失败: {e}")


# def DID_accept():
#     global socket_conn, pending_conn, receiver_did, server_socket
#
#     try:
#         # 如果有已完成协商的连接，直接返回
#         if pending_conn:
#             conn = pending_conn
#             addr = conn.getpeername()  # 获取客户端地址
#             print(f"已获取已协商的连接：{addr}")
#             # 将全局通信socket切换为该连接
#             socket_conn = conn
#             pending_conn = None
#             return conn, addr
#         else:
#             # 没有暂存连接，等待新连接
#             print("等待新连接...")
#
#             new_conn, new_addr = socket_conn.accept()
#             print(f"收到新连接请求：{new_addr}")
#
#             pending_conn = new_conn
#             socket_conn = new_conn
#             return new_conn, new_addr
#
#     except Exception as e:
#         print(f"DID接收连接失败: {e}")
#         return None, None

def DID_accept():

    global server_socket, socket_conn

    try:

        conn, addr = server_socket.accept()
        socket_conn = conn
        return conn, addr
    except Exception as e:
        print(f"DID接收连接失败: {e}")
        return None